import numpy as np
import matplotlib.pyplot as plt

data = []
path = "./resources/data1.txt"


def loadData(path):
    with open(path, 'r') as fp:
        for line in fp.readlines():
            tempList = line.strip().split(',')
            data.append([float(tempList[1]), float(tempList[2]), int(tempList[3])])
    fp.close()


def showData(p, clustering, t):
    color = ['blue', 'yellow', 'green', 'orange', 'purple']
    size = [25, 100]

    plt.title('LVQ')
    plt.xlabel('density')
    plt.ylabel('sweetness')
    for i in range(len(clustering)):
        plotSize = size[t[i]]
        for j in range(len(clustering[i])):
            plt.scatter(clustering[i][j][0], clustering[i][j][1], c=color[i], s=plotSize)

    for i in range(len(p)):
        plt.scatter(p[i][0], p[i][1], s=100, c="red", marker="+")

    plt.show()


def LVQ(n_cluster, max_iter, learning_rate):
    # 根据假定的簇类数量，初始化原型向量以及对应的标签
    p = []
    t = []
    for _ in range(n_cluster):
        index = np.random.randint(len(data))
        p.append(data[index][:-1])
        t.append(data[index][-1])
    p_array = np.array(p)
    data_array = np.array(data)
    current_iter = 0
    while current_iter <= max_iter:
        current_iter += 1
        # 随机选择一个样本
        randomIndex = np.random.randint(len(data))
        # 计算该样本与原型向量的最小距离的下标
        minDistIndex = np.linalg.norm(data_array[randomIndex][:-1]-p_array, axis=1).argmin()
        # 更新原型向量
        if data_array[randomIndex][-1] == t[minDistIndex]:
            p_array[minDistIndex] += learning_rate * (data_array[randomIndex][:-1]-p_array[minDistIndex])
        else:
            p_array[minDistIndex] -= learning_rate * (data_array[randomIndex][:-1] - p_array[minDistIndex])
    c = {}
    # 初始化簇类
    for i in range(n_cluster):
        c[i] = []
    # 将样本划分到与其距离最近的原型向量所代表的簇类
    for i in range(len(data_array)):
        finalMinDistIndex = np.linalg.norm(data_array[i][:-1]-p_array, axis=1).argmin()
        c[finalMinDistIndex].append(data_array[i][:-1])
    return p_array, c, t


if __name__ == '__main__':
    loadData(path)
    p, clustering, t = LVQ(5, 400, 0.1)
    print(p)
    print(clustering)
    showData(p, clustering, t)